import torch
from flask import Flask, render_template, jsonify, request
from app import create_app
#导入fast目录下的ft_model方法
from app.model_train.fast.ft_model import FastTextModel
from app.model_train.bert.model_train_predict import BertModelTrain
from app.model_train.bert.knowledge_distillation import KnowledgeDistillation
from app.model_train.bert.config import Config
from app.model_train.bert.textCNN import TextCnnModel

app = create_app()
app.config['TEMPLATES_AUTO_RELOAD'] = True
app.jinja_env.cache = None

@app.route('/')
def hello_world():  # put application's code here
    return render_template("index.html")

@app.route('/ft_model', methods=['POST'])
def ft_model():
    #实例化对象
    fast_model = FastTextModel()
    # #获取训练集数据
    # fast_model.train_data_save()
    # #获取测试集数据
    # fast_model.test_data_save()
    # #获取测试集数据
    # fast_model.dev_data_save()
    # #模型训练与测试
    # result = fast_model.fast_model_train()
    # print(result)
    # #模型优化
    # fast_model.fast_model_optimization()

    #模型预测
    #接受数据
    text = request.get_json()
    result = fast_model.model_predict(text['text'])
    return jsonify({"success": 1, "result": result})

@app.route("/bert_model")
def bert_model():
    #实例化对象
    bert_model_trainer = BertModelTrain()
    # #模型训练
    # bert_model_trainer.model_application("train")
    # #模型测试
    # bert_model_trainer.model_application("test")
    #模型量化
    bert_model_trainer.model_application("quantization")
    return "ok"

@app.route("/bert_inference", methods=['POST'])
def bert_inference():
    #接受文本参数
    text = request.get_json()

    fast_model = FastTextModel()
    bert_model1 = BertModelTrain()
    knn_model = BertModelTrain()

    fast_result = fast_model.model_predict(text['text'])
    bert_result = bert_model1.model_application("inference", text["text"])
    knn_result = knn_model.textcnn_unified_predict(text["text"])

    return jsonify({"success": 200, "result": f"BERT模型推理结果:{bert_result},随机森林模型推理结果:{fast_result},cnn模型推理结果:{knn_result}"})

@app.route("/kd")
def knowledge_distillation():
    kd = KnowledgeDistillation()
    kd.model_distillation("kd_train")


if __name__ == '__main__':
    # app.run()
    #修改端口号
    app.run(host='0.0.0.0', port=8090, debug=True)